import json
import cv2
import numpy as np

def solver_rigid(pts_3d, pts_2d, camera_matrix):
    # pts_3d  Nx3
    # pts_2d  Nx2
    # camera_matrix 4x4
    dist_coeffs = np.zeros((4, 1))
    pts_3d = pts_3d.copy()
    pts_2d = pts_2d.copy()
    # print(pts_3d.shape, pts_3d.dtype, pts_2d.shape, pts_2d.dtype)
    success, rotation_vector, translation_vector = cv2.solvePnP(pts_3d, pts_2d, camera_matrix, dist_coeffs, flags=0)
    assert success
    R, _ = cv2.Rodrigues(rotation_vector)
    R = R.T
    R[:, 1:3] *= -1
    T = translation_vector.flatten()
    T[1:] *= -1

    return R, T

img = cv2.imread("wcpa/1615021141863_ar.jpg")
img_h, img_w, _ = img.shape

pts68 = np.loadtxt('wcpa/1615021141863_68landmarks.txt', dtype=np.int32)
txt_path = 'resources/projection_matrix.txt'

M_proj = np.loadtxt(txt_path, dtype=np.float32)
print(M_proj.shape) # M_proj shape is (4, 4)

M = np.load("wcpa/1615021141863_info.npz")
verts3d, R_t = M['verts'], M['R_t']
print(R_t.shape) # R_t shape is (4, 4)

npy_path = 'resources/kpt_ind.npy'
kpt_ind = np.load(npy_path)
print(kpt_ind.shape)

# print(verts3d, )

ones = np.ones([verts3d.shape[0], 1])
verts_homo = np.concatenate([verts3d, ones], axis=1)

M1 = np.array([
    [img_w / 2, 0, 0, 0],
    [0, img_h / 2, 0, 0],
    [0, 0, 1, 0],
    [img_w / 2, img_h / 2, 0, 1]
])

verts = verts_homo @ R_t @ M_proj @ M1
verts_cp = verts.copy()
w_ = verts[:, [3]]
verts = verts / w_

# image space: →+x，↓+y
points2d = verts[:, :2]
# 计算后的投影坐标
points2d[:, 1] = img_h - points2d[:, 1]

# print(points2d)
temp1 = img.copy()
for p in points2d:
    cv2.circle(temp1, (int(p[0]), int(p[1])), radius=1, color=(0, 255, 0), thickness=-1)

for p in points2d[kpt_ind]:
    cv2.circle(temp1, (int(p[0]), int(p[1])), radius=1, color=(0, 0, 255), thickness=-1)

print('R: ', R_t)

# camera_matrix = M_proj[:3, :3].T
# r, t = solver_rigid(verts3d, points2d, camera_matrix)
# print(r, t)


cv2.imshow("temp1", temp1)
cv2.waitKey(0)